python3 train.py \
    --task=mnist \
    --dataset-root=/private/benchmarks \
    --ckpt-path=./mnist_ckpt